#!/bin/bash

#----------------------------------------------------
# Default parameters
#----------------------------------------------------
app_path=/home/gop/zelixir/Apps/
foldseek=$app_path/foldseek/bin/foldseek
usalign=$app_path/USalign/bin/USalign
data_path=/home/gop/zelixir/db_test/
n_model=3
tm_index=3
tm_threshold=0.3
print_help=0
work_dir='output'

#----------------------------------------------------
# for test and debug
#----------------------------------------------------
Classification=Bacteria # Bacteria  Phage  ProPhage All
Species=SpeciesTaxid_357276
input_pdb=1UBQ.pdb

#----------------------------------------------------
# parse parameters
#----------------------------------------------------

POSITIONAL=()
while [[ $# -gt 0 ]]; do
 key="$1"
 case $key in
    -h|--help)
     print_help=1
     shift # past argument
     ;;
    -q|--query)
     input_pdb=($2)
     shift # past argument
     shift # past value
     ;;
    -o|--output)
     work_dir=($2)
     shift # past argument
     shift # past value
     ;;
    -c|--class)
     Classification="$2"
     shift # past argument
     shift # past value
     ;;
    -s|--species)
     Species=($2)
     shift # past argument
     shift # past value
     ;;
    --database-root)
     data_path="$2"
     shift # past argument
     shift # past value
     ;;
    --tm-threshold)
     tm_threshold=($2)
     shift # past argument
     shift # past value
     ;;
    --n-best-hits)
     n_model=($2)
     shift # past argument
     shift # past value
     ;;
    --sort-criterion)
     tm_index=($2)
     shift # past argument
     shift # past value
     ;;
    --foldseek)
     foldseek="$2"
     shift # past argument
     shift # past value
     ;;
    --usalign)
     usalign="$2"
     shift # past argument
     shift # past value
     ;;
     *)
   POSITIONAL+=("$1") # save it in an array for later
     shift # past argument
     ;;
 esac
done

#----------------------------------------------------
# print help 
#----------------------------------------------------

help_info(){
    echo ""
    echo 'Usage: DLDB_Search.sh -q query.pdb -c Bacteria -s SpeciesTaxid_357276 -o output_path'
    echo ""
    echo '            -q|--query '
    exit 1
}

if [[ $print_help -eq 1 ]]; then
    help_info
fi

#----------------------------------------------------
# prepare 
#----------------------------------------------------

Species_path=$data_path/$Classification/$Species
pdb_path=$Species_path/pdb/
csv_path=$Species_path/$Species.csv
foldseek_db=$Species_path/foldseek/fsDB

mkdir -p $work_dir
cp $input_pdb $work_dir
pdb_base=`basename $input_pdb`
cd  $work_dir

#----------------------------------------------------
# do foldseek  
#----------------------------------------------------

format=query,target,fident,pident,alnlen,qcov,qlen,tcov,tlen,evalue,alntmscore,lddt,bits,prob

$foldseek easy-search $pdb_base $foldseek_db foldseek.out tmpFolder --exhaustive-search 1 --tmscore-threshold $tm_threshold --format-output $format

#----------------------------------------------------
# do usalign
#----------------------------------------------------

awk '{print $2}' foldseek.out > foldseek.hits
$usalign $pdb_base -dir2 $pdb_path foldseek.hits  -split 0 -outfmt 2 > usalign.out

#----------------------------------------------------
# give models of the best hits
#----------------------------------------------------

let tm_index=tm_index+1
grep "^#" -v usalign.out | awk '{print $3}' > TM1.out
grep "^#" -v usalign.out | awk '{print $4}' > TM2.out
paste TM1.out TM2.out | awk '{if ($1<$2) {print $1} else {print $2}}' > TMmin.out
paste TM1.out TM2.out | awk '{if ($1<$2) {print $2} else {print $1}}' > TMmax.out
paste TM1.out TM2.out | awk '{print ($1+$2)*0.5}' > TMmean.out
paste TM1.out TM2.out TMmin.out TMmax.out TMmean.out TMmin.out | nl > TMall.out
rm TM1.out TM2.out TMmin.out TMmax.out TMmean.out

cat TMall.out | sort -n -k $tm_index -r | awk '{if ($a>=0.5) {print $0}}' a=$tm_index > TMsort.out
n_hit=`cat TMsort.out | wc -l`

if [[ $n_hit -eq 0 ]] && [[ $tm_index -eq 7 ]]; then
    tm_index=5
    cat TMall.out | sort -n -k $tm_index -r | awk '{if ($a>=0.5) {print $0}}' a=$tm_index > TMsort.out
    n_hit=`cat TMsort.out | wc -l`
fi

head -n 1 usalign.out | sed 's/#//g' > usalign_hits.csv

touch pdb.list
rm pdb.list
mkdir -p best_hit_models

if [[ $n_hit -eq 0 ]]; then
    echo "no hit" >> usalign_hits.csv
else
    awk '{print $1}' TMsort.out > used_index.dat
    echo "rank" > rank.list
    ii=1
    for n in `cat used_index.dat`
    do
        grep "^#" -v usalign.out | head -n $n | tail -n 1  >> usalign_hits.csv
        echo "rank"$ii >> rank.list	
        pdb=`grep "^#" -v usalign.out | head -n $n | tail -n 1 | awk '{print $2}' | sed 's/\///g'`
	echo "$pdb" >> pdb.list
	let ii=ii+1
    done
    paste rank.list usalign_hits.csv | tr '\t' ',' > usalign_hits.csv_added
    mv usalign_hits.csv_added usalign_hits.csv
    

    if [[ $n_hit -lt $n_model ]]; then
        n_model=$n_hit
    fi
    
    for n in `head -n $n_model used_index.dat`
    do
        pdb=`grep "^#" -v usalign.out | head -n $n | tail -n 1 | awk '{print $2}' | sed 's/\///g'`
	cp $pdb_path/$pdb ./best_hit_models
    done
fi

#----------------------------------------------------
# give models of the best hits
#----------------------------------------------------

header=`head -n 1  $csv_path`
echo 'rank,'$header > annotation.csv

ii=1
for pdb in `cat pdb.list | sed "s/.pdb$//g"`
do
   ann=`grep "$Species/pdb/$pdb" $csv_path`
   if [[ $ann == '' ]]; then
       ann=$Species/pdb/$pdb
   fi
   echo 'rank'$ii,$ann >> annotation.csv
   let ii=ii+1
done

query_base=`echo $pdb_base| sed "s/.pdb$//g"`

sub_foler=c${Classification}_s${Species}_q${query_base}
mkdir -p $sub_foler

mv best_hit_models annotation.csv usalign_hits.csv $sub_foler
mv $pdb_base $sub_foler

mv pdb.list rank.list foldseek.hits foldseek.out usalign.out TMsort.out TMall.out used_index.dat tmpFolder
rm -rf tmpFolder

mkdir -p results
mv $sub_foler results
